Folder structure:¶Batchsize needs to be changed based on your system specifications.Steps:¶For training, each image is converted in to binary image, where 0's represent background, 1's represent wheat spikes Encoding is done using bounding boxes coordinates present in the Train csv.
images are resized to reduced computation on CPU.
import cv2
import math
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
### Scikit learn libraries ###
from skimage.color import rgb2gray
from skimage.feature import blob_dog
from keras.models import model_from_json
### Pyspark libraries ###
import os,gc
import shutil
import pyspark
import itertools
from operator import add
from pyspark.sql.types import *
from pyspark import SparkContext
from pyspark.sql import DataFrame
from pyspark.sql.functions import udf
from pyspark.ml.image import ImageSchema
file = pd.read_csv("train_csv/train.csv") # This file contains bbox co-ordinates of wheat spikes for each image.
# Train csv has multiple rows for each image depending upon number of wheat spikes in each image.
ids = [file.iloc[i][0] for i in range(file.shape[0])]
# Factor by which we will resize each image as each raw image has size 1024*1024, which will be too large for network to train on to.
factor = 2; size = 512; batch_size = 500; x_axis_gaps = 5; coords_dist_th = 15; N =2999 ;N_ = N+30
sample = os.listdir("train")[0:3]
for img_name in sample:
image = cv2.imread('train/'+ img_name)
# Get image ids for each image to access the segmentation co-ordinates
image_ids = [l for l,val in enumerate(ids) if val == str(img_name.split(".")[0])]
list_of_coords = [[int(float(val)) for val in file.iloc[v][3][1:len(file.iloc[v][3])-1].split(",")] for v in image_ids]
for l in list_of_coords:
image_annotated = cv2.rectangle(image,(int(l[0]),int(l[1])),(int((l[0]+l[2])),int((l[1]+l[3]))),(255,0,0),4)
fig, axes = plt.subplots(1, 2, figsize=(10, 5))#, sharex=True, sharey=True)
ax = axes.ravel()
image = cv2.imread('train/'+ img_name)
ax[0].imshow(cv2.cvtColor(image,cv2.COLOR_RGB2BGR))
ax[1].imshow(cv2.cvtColor(image_annotated,cv2.COLOR_RGB2BGR))
plt.show()
Let's get started with Data Preparation ................¶if not os.path.exists('train_annotated'):
os.makedirs('train_annotated')
from data_prep_module import data_prep
Setting up Spark Context¶Note, I made a few chnages in the backend code for Pyspark UDF to work here. In case you face challanges, feel free to comment/connect.
def run_spark():
sc=SparkContext(master="local[15]") # number[15] can be changed based on your system specification.
# print(sc.binaryFiles("spark_temp/*.png"))
image_df = sc.binaryFiles("spark_temp") # Reading images into Spark context
pyspark.sql.udf.UDFRegistration.register(name="data_prep", f = data_prep, returnType=StringType()) #registering UDF
# such that Spark context recongises the function used for data preparation.
job = [data_prep(ids,factor,file,size,x[0].split("/")[-1]) for idx,x in enumerate(image_df.take(batch_size))]
shutil.rmtree('spark_temp') #remove the folder once 'one batch' is complete to avoid Spark remembering
# indices for images it is done with.
sc.stop() # stop the spark context else Spark will have unnecessary information cached whcih we do not require anymore.
gc.collect() # remove any other cache which memory might have been holding.
input_list = os.listdir("train")[0:N]
input_seq = sorted(list(set([val for val in range(0, len(input_list),batch_size)]+ [len(input_list)])))
def image_transfer(value):
if not os.path.exists('spark_temp'):
os.makedirs('spark_temp')
image = cv2.imread('train/'+value)
cv2.imwrite("spark_temp/" + value.split(".")[0]+".jpg",cv2.resize(image,(size,size)))
for val in range(len(input_seq)-1):
img_batch = input_list[input_seq[val]:input_seq[val]+batch_size]
[image_transfer(val) for val in img_batch]
run_spark()
Lets' do a bit checking on our data.
if not os.path.exists('train_temp'):
os.makedirs('train_temp')
def image_transfer_(value):
cv2.imwrite("train_temp/" + value.split(".")[0]+".png",cv2.resize(cv2.imread('train/'+value.split(".")[0]+".jpg"),(size,size)))
first_3000 = os.listdir("train_annotated")
task = [image_transfer_(val) for val in first_3000]
error_list = [val if np.max(cv2.imread('train_annotated/'+val)) > 1 else None for val in os.listdir("train_annotated")]
error_list = [val for val in error_list if val != None]
if len(error_list) > 0:
print(error_list)
[os.remove("train_temp/"+val.split(".")[0]+".jpg") for val in error_list]
[os.remove("train_annotated/"+val) for val in error_list]
# import tensorflow.compat.v1 as tensorflow
from keras_segmentation.models.unet import vgg_unet
model = vgg_unet(n_classes=2, input_height=size, input_width=size)
model.train(
train_images = "train_temp/",#train
train_annotations = "train_annotated/",n_classes = 2,epochs=10,steps_per_epoch=5,#annotations_prepped_train_v3
)
#saving model to disk
from keras.models import model_from_json
model_json = model.to_json()
with open("model.json", "w") as json_file:
json_file.write(model_json)
model.save_weights("model.h5")
for val in os.listdir("test/"):
cv2.imwrite("test/"+val.split(".")[0]+".jpg",cv2.resize(cv2.imread('test/'+val.split(".")[0]+".jpg"),(size,size)))
import warnings
warnings.simplefilter("ignore")
accuracy = []; submission = {"image_id" : [], "width" : [], "height" : [],"bbox" : [], "Prob":[]}
for idx,img_name in enumerate(os.listdir("test/")[0:5]):
if not os.path.exists('segmentation'):
os.makedirs('segmentation')
if not os.path.exists('Output'):
os.makedirs('Output')
out = model.predict_segmentation(
inp="test/"+ img_name,
out_fname="segmentation/" + img_name.split(".")[0] + ".png"
)
out_image = cv2.imread('segmentation/'+ img_name.split(".")[0] + ".png")
_image = cv2.imread('test/'+ img_name.split(".")[0] + ".jpg")
## Step 4:Blob detection
image_gray = rgb2gray(out_image)
blobs_dog = blob_dog(image_gray, max_sigma=30, threshold=.10)
try:
blobs_dog[:, 2] = blobs_dog[:, 2] * math.sqrt(2)
submission["image_id"].append(img_name.split(".")[0])
submission["width"].append(size); submission["height"].append(size)
submission["bbox"].append([str([int(blob[0]), int(blob[1]),int(blob[2]), int(blob[2])]) for blob in blobs_dog])
submission["Prob"].append(.50)
except:
submission["image_id"].append(img_name.split(".")[0])
submission["width"].append(size); submission["height"].append(size)
submission["bbox"].append([])
submission["Prob"].append(1)
blob_list_sorted = sorted([[int(blob[0]), int(blob[1]),int(blob[2])] for blob in blobs_dog])
len_val = len(blob_list_sorted)
for i,val in enumerate(range(0,len(blob_list_sorted))):
for j in range(i+1,len(blob_list_sorted)-2):
if max(i,j) < len_val and max([abs(np.diff(x)) for x in zip(blob_list_sorted[j][0:2], blob_list_sorted[i][0:2])]) < coords_dist_th:
blob_list_sorted = blob_list_sorted+[[min(x) for x in zip(blob_list_sorted[j], blob_list_sorted[i])][0:2]+[blob_list_sorted[j][2]+blob_list_sorted[i][2]]]
del blob_list_sorted[i]
del blob_list_sorted[j]
len_val = len(blob_list_sorted)
try:
for blob in blobs_dog:
y, x, r = blob
Iou = cv2.rectangle(_image,(int(x),int(y)),(int((x+r)),int((y+r))),(255,0,0),2)
# print(Iou)
cv2.imwrite("Output/" + img_name.split(".")[0] + ".png",Iou)
fig, axes = plt.subplots(1, 2, figsize=(12, 6))#, sharex=True, sharey=True)
ax = axes.ravel()
ax[0].imshow(cv2.cvtColor(cv2.imread('test/'+ img_name.split(".")[0] + ".jpg"),cv2.COLOR_RGB2BGR))
ax[1].imshow(Iou)
plt.show()
except:pass
pd.DataFrame(submission).to_csv("submission.csv")
pd.DataFrame(submission).to_csv("submission.csv")
Results are looking alryt! Let's test on few more images..
if not os.path.exists('train_50'):
os.makedirs('train_50')
for val in os.listdir("train/")[N:N_]:
cv2.imwrite("train_50/"+val.split(".")[0]+".png",cv2.resize(cv2.imread('train/'+val.split(".")[0]+".jpg"),(size,size)))
warnings.simplefilter("ignore")
for idx,img_name in enumerate(os.listdir("train_50")):
if not os.path.exists('segmentation'):
os.makedirs('segmentation')
out = model.predict_segmentation(
inp="train_50/"+ img_name,
out_fname="segmentation/" + img_name.split(".")[0] + ".png"
)
out_image = cv2.imread('segmentation/'+ img_name.split(".")[0] + ".png")
_image = cv2.imread('train_50/'+ img_name.split(".")[0] + ".png")
## Step 4:Blob detection
image_gray = rgb2gray(out_image)
blobs_dog = blob_dog(image_gray, max_sigma=30, threshold=.10)
blob_list_sorted = sorted([[int(blob[0]), int(blob[1]),int(blob[2])] for blob in blobs_dog])
len_val = len(blob_list_sorted)
for i,val in enumerate(range(0,len(blob_list_sorted))):
for j in range(i+1,len(blob_list_sorted)-2):
if max(i,j) < len_val and max([abs(np.diff(x)) for x in zip(blob_list_sorted[j][0:2], blob_list_sorted[i][0:2])]) < coords_dist_th:
blob_list_sorted = blob_list_sorted+[[min(x) for x in zip(blob_list_sorted[j], blob_list_sorted[i])][0:2]+[blob_list_sorted[j][2]+blob_list_sorted[i][2]]]
del blob_list_sorted[i]
del blob_list_sorted[j]
len_val = len(blob_list_sorted)
try:
for blob in blobs_dog:
y, x, r = blob
Iou = cv2.rectangle(_image,(int(x),int(y)),(int((x+r)),int((y+r))),(255,0,0),2)
fig, axes = plt.subplots(1, 2, figsize=(12, 6))
ax = axes.ravel()
ax[0].imshow(cv2.cvtColor(cv2.imread('train_50/'+ img_name.split(".")[0] + ".png"),cv2.COLOR_RGB2BGR))
ax[1].imshow(Iou)
plt.show()
except:pass
To all my fellow Data Scientists, non privileged ones, privileged ones, in the industry because of relatives, or by yourself... Happy coding ;)
This work took me a month time to identify right pre-processing technique which worked for all (almost) images. if you notice, results are pretty satisfactory in terms of algorithm picking all the wheat heads in an image despite the varied size of the wheat heads, different light conditions and orientation.
I agree, there is scope for improvement (which I will take up later in sometime..) but I am equally proud that the solution is able to seperate wheat heads from leaves which are of same color and at times, same shape and size. and top that solution is able to discard the background which is soil, dry leaves and etc..
Few Applications of this exercise:
Appreciate the effort.., do critic, share feedback. You all are most welcome.